package cn.cscamp.middleware.loadbalancer;

import cn.cscamp.middleware.loadbalancer.config.LoadbalancerService;
import cn.cscamp.middleware.loadbalancer.config.LoadbalancerServiceConfig;
import cn.cscamp.middleware.loadbalancer.exception.LoadbalancerException;
import cn.cscamp.middleware.loadbalancer.strategy.RandomLoadbalancer;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.HttpStatus;
import org.springframework.web.reactive.function.client.ClientRequest;
import org.springframework.web.reactive.function.client.ClientResponse;
import org.springframework.web.reactive.function.client.ExchangeFilterFunction;
import org.springframework.web.reactive.function.client.ExchangeFunction;
import org.springframework.web.util.UriComponentsBuilder;
import reactor.core.publisher.Mono;
import reactor.retry.Retry;

import java.net.URI;
import java.util.List;

@SuppressWarnings("all")
public class CustomRetryableLoadBalancerExchangeFilterFunction implements ExchangeFilterFunction {

    private static final Logger LOGGER = LoggerFactory.getLogger(CustomRetryableLoadBalancerExchangeFilterFunction.class);

    private final LoadbalancerServiceConfig loadbalancerServiceConfig;

    private final List<LoadbalancerService> loadbalancerServices;

    public CustomRetryableLoadBalancerExchangeFilterFunction(LoadbalancerServiceConfig loadbalancerServiceConfig) {
        this.loadbalancerServiceConfig = loadbalancerServiceConfig;
        this.loadbalancerServices = loadbalancerServiceConfig.getLoadbalancerServices();
    }

    @Override
    public Mono<ClientResponse> filter(ClientRequest request, ExchangeFunction next) throws LoadbalancerException {
        URI originalUrl = request.url();
        String serviceId = originalUrl.getHost();

        if (serviceId == null || loadbalancerServices.stream()
                        .noneMatch(service -> serviceId.equals(service.getServiceId()))
        ) {
            String message = String.format("Request URI does not contain a valid hostname: %s", originalUrl);
            if (LOGGER.isWarnEnabled()) {
                LOGGER.warn(message);
            }
            return Mono.just(ClientResponse.create(HttpStatus.BAD_REQUEST).body(message).build());
        }

        //匹配出对应的serviceAddress列表
        List<String> serviceAddressList = loadbalancerServices.stream()
                .filter(service -> serviceId.equals(service.getServiceId()))
                .findFirst()
                .orElseThrow(() -> new RuntimeException("cannot get valid service address"))
                .getServiceAddress();

        RandomLoadbalancer randomLoadbalancer = new RandomLoadbalancer(serviceAddressList);

        return Mono.defer(() -> {
            ClientRequest newRequest = buildClientRequest(request, randomLoadbalancer.chooseOne());
            return next.exchange(newRequest)
                    .doOnSuccess(clientResponse -> {
                        //ignored
                    })
                    .doOnError(throwable -> {
                        randomLoadbalancer.onError();
                        LOGGER.warn("request url: {} failed.", newRequest.url().toString());
                        throw new LoadbalancerException("try to connect all urls, cannot find any valid url.");
                    });
        }).retryWhen(Retry.onlyIf(x -> x.exception() instanceof LoadbalancerException)
                .retryMax(randomLoadbalancer.getRetry()));
    }

    static ClientRequest buildClientRequest(ClientRequest request, String serviceAddress) {
        URI originalUrl = request.url();
        return ClientRequest.create(request.method(), reconstructURI(serviceAddress, originalUrl))
                .headers(headers -> headers.addAll(request.headers())).cookies(cookies -> cookies.addAll(request.cookies()))
                .attributes(attributes -> attributes.putAll(request.attributes()))
                .body(request.body())
                .build();
    }

    static URI reconstructURI(String serviceAddress, URI original) {
        String[] split = serviceAddress.split(":");
        String host = split[0];
        int port = Integer.valueOf(split[1]);
        return UriComponentsBuilder.fromUri(original).host(host).port(port).build().toUri();
    }
}
